import io
import json
import os
import shutil
import socket
import tempfile
import zipfile

import pandas as pd
from fastapi import FastAPI, UploadFile, HTTPException, BackgroundTasks
from fastapi.params import File
from fastapi.responses import FileResponse
from starlette.responses import JSONResponse, StreamingResponse

from system.apiTools.utils import (delete_files_in_folder, remove_file, remove_file1)
from system.predict_csv import start_application
# from system.tools.infer.predict_system import start_application

app = FastAPI()
UPLOAD_DIRECTORY = "test"
directory = "inference_results"
# 文件夹路径
FILE_DIRECTORY = 'excel_results'


@app.post("/imageDownload")
async def imageDownload(background_tasks: BackgroundTasks, file: UploadFile = File(...)):

    base_name, ext = os.path.splitext(file.filename)
    file_name = base_name + ".xlsx"
    # 确保文件存在
    file_path = os.path.join(FILE_DIRECTORY, file_name)
    if not os.path.exists(file_path) or not file_path.endswith('.xlsx'):
        resMap = {"code": 200, "data": "文件不存在"}
        return JSONResponse(content=resMap, status_code=200)

    # 添加删除文件的后台任务
    background_tasks.add_task(remove_file, file_path)

    return FileResponse(path=file_path, filename=file_name)


@app.post("/pdfDownload")
async def pdfDownload(background_tasks: BackgroundTasks, file: UploadFile = File(...)):
    # 允许的文件类型
    allowed_file_types = ["application/pdf"]

    # 检查文件类型
    if file.content_type not in allowed_file_types:
        raise HTTPException(status_code=400, detail="Invalid file type")

    # 构造文件路径
    file_path = os.path.join(UPLOAD_DIRECTORY, file.filename)
    print(file_path)

    # 保存文件到指定目录
    with open(file_path, "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)

    try:
        file_name = file.filename
        identificationResults = start_application(write2excel=True, image_dir=file_path)
        res_list = []
        for res in identificationResults:
            res = res[len(file_name) + 2:]
            res_list.append(json.loads(res))
        data_map = {"fileName": file_name, "data": res_list}
        resMap = {"code": 200, "data": data_map}
    except Exception as e:
        print(e)
        resMap = {"code": 403, "data": "识别失败"}

    os.remove(file_path)

    base_name, ext = os.path.splitext(file.filename)
    # 指定的文件夹路径
    folder_path = FILE_DIRECTORY
    # 获取文件夹下所有文件
    files = os.listdir(folder_path)
    # 筛选以指定前缀开头的excel文件
    excel_files = [file for file in files if file.startswith(base_name) and file.endswith('.xlsx')]
    print(excel_files)

    # 检查是否找到了符合条件的文件
    if not excel_files:
        raise HTTPException(status_code=404, detail="No files found with the specified prefix.")

    # 创建一个字节流对象
    zip_buffer = io.BytesIO()
    with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zipf:
        for file in excel_files:
            file_path = os.path.join(folder_path, file)
            zipf.write(file_path, arcname=os.path.basename(file))

    # 重置字节流指针到起始位置
    zip_buffer.seek(0)

    # # 添加文件删除任务到后台任务
    # background_tasks.add_task(remove_file, temp_zip_file.name)
    # # 添加删除文件的后台任务
    background_tasks.add_task(delete_files_in_folder, FILE_DIRECTORY, excel_files)

    return StreamingResponse(zip_buffer, media_type='application/x-zip-compressed',
                             headers={"Content-Disposition": f"attachment; filename={base_name}_files.zip"})


@app.post("/imageDetect")
async def upload_file(background_tasks: BackgroundTasks, file: UploadFile = File(...)):
    # 检查目录是否存在
    if os.path.exists(directory) and os.path.isdir(directory):
        print(f"目录 {directory} 存在。")
        shutil.rmtree(directory)
    os.mkdir(directory)
    # 检查目录是否存在
    if os.path.exists(UPLOAD_DIRECTORY) and os.path.isdir(UPLOAD_DIRECTORY):
        print(f"目录 {UPLOAD_DIRECTORY} 存在。")
        shutil.rmtree(UPLOAD_DIRECTORY)
    os.mkdir(UPLOAD_DIRECTORY)
    # 允许的文件类型
    allowed_file_types = ["image/jpeg", "image/png"]
    # 检查文件类型
    if file.content_type not in allowed_file_types:
        resMap = {"code": 403, "data": "文件类别错误"}
        return JSONResponse(content=resMap, status_code=200)
    # 构造文件路径
    file_path = os.path.join(UPLOAD_DIRECTORY, file.filename)
    # 保存文件到指定目录
    with open(file_path, "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)
    try:
        file_name = file.filename
        identificationResults = start_application(write2excel=True, image_dir=file_path)
        identificationResults[0] = identificationResults[0][len(file_name) + 2:-2]
        # print("--------------------")
        # print(type(identificationResults))
        # print("--------------------")
        # print(identificationResults[0])
        res_json = json.loads("[" + identificationResults[0] + "]")
        # res_json = json.loads(identificationResults[0])
        data_map = {"fileName": file_name, "data": res_json}
        resMap = {"code": 200, "data": data_map}
    except Exception as e:
        print(e)
        resMap = {"code": 403, "data": "识别失败"}
        # 添加删除文件的后台任务
    background_tasks.add_task(remove_file1, file_path, 5)
    # os.remove(file_path)

    base_name, ext = os.path.splitext(file.filename)
    file_name = base_name + ".xlsx"
    # 确保文件存在
    file_path = os.path.join(FILE_DIRECTORY, file_name)
    if not os.path.exists(file_path) or not file_path.endswith('.xlsx'):
        resMap = {"code": 200, "data": "excel文件不存在"}
        return JSONResponse(content=resMap, status_code=200)

    # 添加删除文件的后台任务
    background_tasks.add_task(remove_file1, file_path, 60*60)

    return JSONResponse(content=resMap, status_code=200)


# @app.post("/pdfDetect")
# async def pdf_detect(file: UploadFile = File(...)):
#     # 检查目录是否存在
#     if os.path.exists(directory) and os.path.isdir(directory):
#         print(f"目录 {directory} 存在。")
#         shutil.rmtree(directory)
#     os.mkdir(directory)
#     # 允许的文件类型
#     allowed_file_types = ["application/pdf"]
#
#     # 检查文件类型
#     if file.content_type not in allowed_file_types:
#         raise HTTPException(status_code=400, detail="Invalid file type")
#
#     # 构造文件路径
#     file_path = os.path.join(UPLOAD_DIRECTORY, file.filename)
#
#     # 保存文件到指定目录
#     with open(file_path, "wb") as buffer:
#         shutil.copyfileobj(file.file, buffer)
#
#     try:
#         file_name = file.filename
#         identificationResults = start_application(write2excel=True)
#         res_list = []
#         for res in identificationResults:
#             res = res[len(file_name) + 2:]
#             res_list.append(json.loads(res))
#             # res_list.append(res)
#         data_map = {"fileName": file_name, "data": res_list}
#         resMap = {"code": 200, "data": data_map}
#     except Exception as e:
#         print(e)
#         resMap = {"code": 403, "data": "识别失败"}
#     os.remove(file_path)
#
#     return JSONResponse(content=resMap, status_code=200)


def get_host_ip():
    """
    查询本机ip地址
    """
    ss = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    ss.connect(('8.8.8.8', 80))
    ip = ss.getsockname()[0]
    ss.close()
    return ip


if __name__ == "__main__":
    import uvicorn
    start_application()

    ip = get_host_ip()
    # uvicorn.run(app, host=ip, port=8000, reload=False)
